import sys
import os
import importlib
import pdb
import mopo

def import_fns(path, file, fns_name):
    full_path = os.path.join(path, file)
    import_path = full_path.replace('/', '.')
    module = importlib.import_module(import_path)
    fns = getattr(module, fns_name)
    return fns

def load_all_mask(path, env_name):
    full_path = os.path.join(path, env_name)
    import_path = full_path.replace('/', '.')
    module = importlib.import_module(import_path)
    env_structrues = {}
    for attr in dir(module):
        if attr.endswith("mask"):
            structrue_name = "_".join([env_name] + attr.split("_")[:-1])
            env_structrues[structrue_name] = getattr(module, attr)
    return env_structrues

mopo_path = os.path.split(mopo.__file__)[0]
cwd = "mopo/mask"
cwd_path = os.path.join(mopo_path, "mask")
files = os.listdir(cwd_path)
# remove __init__.py base.py .DS_store
files = filter(lambda x: '__' not in x and x[0] != '.' and "base" not in x, files)
# env.py --> env
env_names = list(map(lambda x: x.replace('.py', ''), files))

# causal_masks = {"{}_causal".format(file): import_fns(cwd, file, "CausalStructure") for file in files}
# oracle_masks = {"{}_oracle".format(file): import_fns(cwd, file, "OracleStructure") for file in files}
# dense_masks = {"{}_dense".format(file): import_fns(cwd, file, "DenseStructure") for file in files}
masks = {}

for env_name in env_names:
    env_masks = load_all_mask(cwd, env_name)
    masks.update(env_masks)

# sys.modules[__name__] = masks


# for name in masks:
#     if name.startswith("hopper"):
#         print(name, masks[name].sum() / (12 * 14))